from collections import defaultdict
from sqlglot import exp

def ast_to_features(ast, depth=0):
    """将AST转换为特征字典（包含结构深度信息）"""
    features = defaultdict(int)
    
    def _walk(node, current_depth):
        # 节点类型+深度作为特征键
        feature_key = f"{type(node).__name__}_depth{current_depth}"
        features[feature_key] += 1
        
        # 对特定节点类型添加额外特征
        if isinstance(node, exp.Column):
            features["Column_total"] += 1
        elif isinstance(node, exp.Join):
            features["Join_total"] += 1
        
        for _, value in node.args.items():
            if isinstance(value, list):
                for child in value:
                    if isinstance(child, exp.Expression):
                        _walk(child, current_depth + 1)
            elif isinstance(value, exp.Expression):
                _walk(value, current_depth + 1)
    
    if ast:
        _walk(ast, 0)
    return features